# 加载环境变量
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv
import os
import json
import mysql.connector

_ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)


# 移除字典中的 None 值
def remove_none_values(data):
    if isinstance(data, dict):
        return {key: remove_none_values(value) for key, value in data.items() if value is not None}
    elif isinstance(data, list):
        return [remove_none_values(item) for item in data if item is not None]
    else:
        return data


# 打印参数。如果参数是有结构的（如字典或列表），则以格式化的 JSON 形式打印；
def print_json(data):
    if (isinstance(data, (list, dict))):
        print(json.dumps(data, indent=4, ensure_ascii=False))
    else:
        print(data)


# 对象序列化
def serialize_json(data):
    if hasattr(data, 'model_dump_json'):
        data = json.loads(data.model_dump_json())

    if (isinstance(data, (list, dict))):
        data = remove_none_values(data)
    return data


def get_ddl_completion(params, model="gpt-3.5-turbo"):
    print("=====GPT请求=====")
    messages = serialize_json(params)
    print_json(messages)

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        tools=[{
            "type": "function",
            "function": {
                "name": "get_table_ddl",
                "description": "根据数据库表名，获取表的DDL",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "table_name": {
                            "type": "array",
                            "description": "数据库表名列表，必须是英文，表名中可能包含下划线",
                        }
                    },
                    "required": ["table_name"],
                }
            }
        }],
    )
    result = serialize_json(response.choices[0].message)
    print("=====GPT回复=====")
    print_json(result)
    return result


def get_sql_completion(params, model="gpt-3.5-turbo"):
    print("=====GPT请求=====")
    messages = serialize_json(params)
    print_json(messages)

    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        tools=[{
            # 摘自 OpenAI 官方示例 https://github.com/openai/openai-cookbook/blob/main/examples/How_to_call_functions_with_chat_models.ipynb
            "type": "function",
            "function": {
                "name": "ask_database",
                "description": "Use this function to answer user questions about business. \
                                Output should be a fully formed SQL query.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": f"""
                                SQL query extracting info to answer the user's question.
                                SQL should be written using this database schema:
                                {database_schema_string}
                                The query should be returned in plain text, not in JSON.
                                The query should only contain grammars supported by SQLite.
                                """,
                        }
                    },
                    "required": ["query"],
                }
            }
        }],
    )
    result = serialize_json(response.choices[0].message)
    print("=====GPT回复=====")
    print_json(result)
    return result


# 获取表信息
def get_table_ddl(table_names):
    global cursor
    db_config = {
        'host': os.getenv("MYSQL_HOST"),
        'user': os.getenv("MYSQL_USER"),
        'password': os.getenv("MYSQL_PASSWORD"),
        'database': os.getenv("MYSQL_DATABASE"),
        'port': 3306,
    }
    # 连接到MySQL数据库
    connection = mysql.connector.connect(**db_config)

    try:
        # 存储结果的字符串
        ddl = ""
        # 获取游标
        cursor = connection.cursor(buffered=True)
        # 遍历表信息
        for table_name in table_names:
            # 查询表信息
            cursor.execute(f"SHOW CREATE TABLE {table_name}")
            # 获取查询结果
            result = cursor.fetchone()
            if result:
                table_ddl = result[1] + ";\n"
                ddl += f"{table_ddl}\n"

        return ddl

    finally:
        # 关闭游标和连接
        cursor.close()
        connection.close()


def ask_database(query):
    cursor.execute(query)
    records = cursor.fetchall()
    return records


#  描述数据库表结构
database_schema_string = """
CREATE TABLE `stu_chapter_exercises_answer` (
  `id` int NOT NULL AUTO_INCREMENT COMMENT '主键id',
  `student_id` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NOT NULL COMMENT '学生id',
  `course_id` bigint DEFAULT NULL COMMENT '课程id',
  `chapter_id` bigint DEFAULT NULL COMMENT '课次id',
  `qid` int DEFAULT NULL COMMENT 'tb_chapter_exercises_explain_video表的id',
  `answer` tinyint DEFAULT NULL COMMENT '答案 0=错 1=对 2=半对',
  `create_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '创建人名字',
  `create_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '创建人id',
  `create_time` datetime DEFAULT NULL COMMENT '创建时间',
  `last_modified_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '最后更新人名字',
  `last_modified_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '最后更新人id',
  `last_modified_time` datetime DEFAULT NULL COMMENT '最后更新时间',
  `bid` bigint DEFAULT NULL COMMENT '批次号',
  `is_history` tinyint DEFAULT '0' COMMENT '是否是历史 0=否 1=是',
  `del_status` tinyint NOT NULL DEFAULT '0' COMMENT '是否删除',
  PRIMARY KEY (`id`) USING BTREE,
  KEY `IDX_CREATE_BY` (`create_by`),
  KEY `IDX_STUDENT_ID` (`student_id`),
  KEY `IDX_CREATE_TIME` (`create_time`),
  KEY `IDX_CHAPTER_ID` (`chapter_id`),
  KEY `IDX_LAST_MODIFIED_BY` (`last_modified_by`)
) ENGINE=InnoDB AUTO_INCREMENT=4681 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci ROW_FORMAT=DYNAMIC COMMENT='学生课次小测答案';

CREATE TABLE `stu_homework_exercises_answer` (
  `id` int NOT NULL AUTO_INCREMENT COMMENT '主键id',
  `student_id` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NOT NULL COMMENT '学生id',
  `course_id` bigint DEFAULT NULL COMMENT '课程id',
  `chapter_id` bigint DEFAULT NULL COMMENT '课次id',
  `qid` int DEFAULT NULL COMMENT 'tb_chapter_exercises_homework 的ID',
  `homework_question_number` int DEFAULT NULL COMMENT '作业题号废弃',
  `answer` tinyint DEFAULT NULL COMMENT '答案 0=错 1=对 2=半对',
  `create_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '创建人名字',
  `create_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '创建人id',
  `create_time` datetime DEFAULT NULL COMMENT '创建时间',
  `last_modified_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '最后更新人名字',
  `last_modified_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '最后更新人id',
  `last_modified_time` datetime DEFAULT NULL COMMENT '最后更新时间',
  `bid` bigint DEFAULT NULL COMMENT '日批次号',
  `is_history` tinyint DEFAULT '0' COMMENT '历史标识 0 当前 1历史',
  `del_status` tinyint NOT NULL DEFAULT '0' COMMENT '是否删除',
  PRIMARY KEY (`id`),
  KEY `IDX_CREATE_BY` (`create_by`),
  KEY `IDX_LAST_MODIFIED_BY` (`last_modified_by`),
  KEY `IDX_CHAPTER_ID` (`chapter_id`),
  KEY `IDX_STUDENT_ID` (`student_id`),
  KEY `IDX_IS_HISTORY` (`is_history`),
  KEY `IDX_CREATE_TIME` (`create_time`)
) ENGINE=InnoDB AUTO_INCREMENT=2571 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='学生课次作业题作答记录表';

CREATE TABLE `tb_chapter_exercises_homework` (
  `id` int NOT NULL AUTO_INCREMENT,
  `parent_id` int DEFAULT '0' COMMENT '父级节点',
  `course_id` bigint DEFAULT NULL COMMENT '课程ID',
  `chapter_id` bigint DEFAULT NULL COMMENT '课次ID',
  `question_id` int DEFAULT NULL COMMENT '题库题ID',
  `video_url` varchar(1024) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '视频url',
  `sort` int DEFAULT '0' COMMENT '排序',
  `question_number` varchar(15) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '题号(显示的)',
  `del_status` tinyint DEFAULT '0' COMMENT '删除标识 0=未删 1=删除',
  `play_time` bigint DEFAULT '0' COMMENT '播放时长',
  `create_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `create_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `create_time` datetime DEFAULT NULL COMMENT '创建时间',
  `last_modified_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `last_modified_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `last_modified_time` datetime DEFAULT NULL COMMENT '最后更新时间',
  `video_check` tinyint DEFAULT '0' COMMENT '0=待检查 1=符合标准 2=不符合标准',
  PRIMARY KEY (`id`),
  KEY `IDX_CHAPTER_ID` (`chapter_id`) USING BTREE,
  KEY `IDX_QUESTION_ID` (`question_id`),
  KEY `IDX_COURSE_ID` (`course_id`)
) ENGINE=InnoDB AUTO_INCREMENT=64353 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='课次练习试题信息';

CREATE TABLE `tb_chapter_exercises_explain_video` (
  `id` int NOT NULL AUTO_INCREMENT,
  `parent_id` int DEFAULT '0' COMMENT '父级节点',
  `course_id` bigint DEFAULT NULL COMMENT '课程id',
  `chapter_id` bigint DEFAULT NULL COMMENT '章节id',
  `question_id` int DEFAULT NULL COMMENT '题库题ID',
  `video_url` varchar(1024) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '视频url',
  `sort` int DEFAULT NULL COMMENT '排序',
  `question_number` varchar(15) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '题号(显示的)',
  `homework_question_number` varchar(512) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL COMMENT '关联的作业题号以逗号分割(一般为作业的大体题号)',
  `snapshot` tinyint NOT NULL DEFAULT '1',
  `del_status` tinyint DEFAULT '0' COMMENT '删除标识 0=未删 1=删除',
  `play_time` bigint DEFAULT '0' COMMENT '播放时长',
  `create_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `create_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `create_time` datetime DEFAULT NULL COMMENT '创建时间',
  `last_modified_name` varchar(64) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `last_modified_by` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci DEFAULT NULL,
  `last_modified_time` datetime DEFAULT NULL COMMENT '最后更新时间',
  `video_check` tinyint DEFAULT '0' COMMENT '0=待检查 1=符合标准 2=不符合标准',
  KEY `id` (`id`) USING BTREE,
  KEY `IDX_CHAPTER_ID` (`chapter_id`) USING BTREE,
  KEY `IDX_QUESTION_ID` (`question_id`),
  KEY `IDX_COURSE_ID` (`course_id`)
) ENGINE=InnoDB AUTO_INCREMENT=62864 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci ROW_FORMAT=DYNAMIC COMMENT='课次小测试题信息';
"""

prompt = "统计每月每件商品的销售额"
# prompt = "这星期消费最高的用户是谁？他买了哪些商品？ 每件商品买了几件？花费多少？"
# prompt = "统计指定学员每个科目的错题总数量（包括：小测和练习），做错和半对的答案都算是错题。表的别名要简写"
prompt = "统计学员A做错的小测和练习的试题所对应的题库ID（question_id），做错和半对的答案都算是错题。表的别名要简写"
messages = [
    {"role": "system", "content": "基于 order 表回答用户问题"},
    {"role": "user", "content": prompt}
]
response = get_sql_completion(messages, "gpt-4")

# response = get_sql_completion(messages)

print("=====SQL查询结果=====")
for tool_call in response['tool_calls']:
    print(tool_call['function']['arguments'])




# 调用方法并打印结果
# result = get_table_ddl(["stu_chapter_exercises_answer", "stu_homework_exercises_answer", "tb_chapter_exercises_homework", "tb_chapter_exercises_explain_video"])
# print(result)